package xclient

import (
	"context"
	"io"
	"log"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"

	. "zt/go-study/rpc"
)

// XClient 负载均衡的客户端
type XClient struct {
	d       Discovery
	mode    SelectMode
	opt     *Option
	mu      sync.Mutex         // 互斥锁
	clients map[string]*Client // 为了尽量地复用已经创建好的 Socket 连接，使用 clients 保存创建成功的 Client 实例
}

var _ io.Closer = (*XClient)(nil)

// NewXClient 实例化
func NewXClient(d Discovery, mode SelectMode, opt *Option) *XClient {
	return &XClient{
		d:       d,
		mode:    mode,
		opt:     opt,
		clients: make(map[string]*Client),
	}
}

// Close 关闭已经建立的链接
func (xc *XClient) Close() error {
	xc.mu.Lock()
	defer xc.mu.Unlock()
	for key, client := range xc.clients {
		// 我不知道如何处理错误，就忽略它吧。
		_ = client.Close()
		delete(xc.clients, key)
	}
	return nil
}

// dial 建立客户端链接
func (xc *XClient) dial(rpcAddr string) (*Client, error) {
	xc.mu.Lock()
	defer xc.mu.Unlock()
	client, ok := xc.clients[rpcAddr]
	if ok && !client.IsAvailable() {
		// 获取的客户端已经无效了
		_ = client.Close()
		delete(xc.clients, rpcAddr)
		client = nil
	}

	// 没有返回缓存的 Client，则说明需要创建新的 Client，缓存并返回
	if client == nil {
		var err error
		client, err = XDial(rpcAddr, xc.opt)
		if err != nil {
			return nil, err
		}
		xc.clients[rpcAddr] = client
	}

	return client, nil
}

// call 调用客户端发起请求
func (xc *XClient) call(rpcAddr string, ctx context.Context, serviceMethod string, args, reply interface{}) error {
	client, err := xc.dial(rpcAddr)
	if err != nil {
		return err
	}
	return client.Call(ctx, serviceMethod, args, reply)
}

// Call 对外提供的调用请求的方法
func (xc *XClient) Call(ctx context.Context, serviceMethod string, args, reply interface{}) error {
	rpcAddr, err := xc.d.Get(xc.mode) // 获取需要调用的服务地址信息
	if err != nil {
		return err
	}
	return xc.call(rpcAddr, ctx, serviceMethod, args, reply)
}

// Broadcast 为在discovery中注册的每个服务器调用指定函数
// 将请求广播到所有的服务实例，如果任意一个实例发生错误，则返回其中一个错误；如果调用成功，则返回其中一个的结果。有以下几点需要注意：
// 1. 为了提升性能，请求是并发的。
// 2. 并发情况下需要使用互斥锁保证 error 和 reply 能被正确赋值。
// 3. 借助 context.WithCancel 确保有错误发生时，快速失败。
func (xc *XClient) Broadcast(ctx context.Context, serviceMethod string, args, reply interface{}) error {
	servers, err := xc.d.GetAll()
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	var mu sync.Mutex
	var e error
	replyDone := reply == nil // 如果reply是nil，则不需要设置值
	ctx, cancel := context.WithCancel(ctx)
	for _, rpcAddr := range servers {
		wg.Add(1)
		go func(rpcAddr string) {
			defer wg.Done()
			var clonedReply interface{}
			if reply != nil {
				clonedReply = reflect.New(reflect.ValueOf(reply).Elem().Type()).Interface()
			}
			err := xc.call(rpcAddr, ctx, serviceMethod, args, clonedReply)
			mu.Lock()
			if err != nil && e == nil {
				e = err
				cancel() // 如果有呼叫失败，请取消未完成的呼叫
			}
			if err == nil && !replyDone {
				reflect.ValueOf(reply).Elem().Set(reflect.ValueOf(clonedReply).Elem())
				replyDone = true
			}
			mu.Unlock()
		}(rpcAddr)
	}
	wg.Wait()
	return e
}

// RegistryDiscovery 注册中心的服务发现
type RegistryDiscovery struct {
	*MultiServersDiscovery               // 嵌套了 MultiServersDiscovery，很多能力可以复用
	registry               string        // 注册中心的地址
	timeout                time.Duration // 服务列表的过期时间
	lastUpdate             time.Time     // 最后从注册中心更新服务列表的时间，默认10s过期，即10s之后，需要从注册中心更新新的列表
}

const defaultUpdateTimeout = time.Second * 10

// NewRegistryDiscovery 实例化
func NewRegistryDiscovery(registryAddr string, timeout time.Duration) *RegistryDiscovery {
	if timeout == 0 {
		timeout = defaultUpdateTimeout
	}

	d := &RegistryDiscovery{
		MultiServersDiscovery: NewMultiServersDiscovery(make([]string, 0)),
		registry:              registryAddr,
		timeout:               timeout,
	}
	return d
}

// Update 更新注册中心的服务
func (d *RegistryDiscovery) Update(servers []string) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.servers = servers
	d.lastUpdate = time.Now()
	return nil
}

// Refresh 刷新注册中心
func (d *RegistryDiscovery) Refresh() error {
	d.mu.Lock()
	defer d.mu.Unlock()
	// 如果最后一次更新时间加上超时时间超过了当前时间，则不需要刷新
	if d.lastUpdate.Add(d.timeout).After(time.Now()) {
		return nil
	}
	log.Println("rpc registry: refresh servers from registry", d.registry)
	resp, err := http.Get(d.registry)
	if err != nil {
		log.Println("rpc registry refresh err:", err)
		return err
	}

	// 创建服务
	servers := strings.Split(resp.Header.Get("X-Geerpc-Servers"), ",")
	d.servers = make([]string, 0, len(servers))
	for _, server := range servers {
		if strings.TrimSpace(server) != "" {
			d.servers = append(d.servers, strings.TrimSpace(server))
		}
	}
	// 更新最后一次时间
	d.lastUpdate = time.Now()
	return nil
}

// Get 从注册中心中获取服务实例
func (d *RegistryDiscovery) Get(mode SelectMode) (string, error) {
	if err := d.Refresh(); err != nil {
	    return "", err
	}
	return d.MultiServersDiscovery.Get(mode)
}

// GetAll 从注册中心中获取服务列表
func (d *RegistryDiscovery) GetAll() ([]string, error) {
	if err := d.Refresh(); err != nil {
	    return nil, err
	}
	return d.MultiServersDiscovery.GetAll()
}
